package middlewares

import (
	"net/http"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/juju/ratelimit"
)

//fillInterval 填充令牌速率 cap 令牌桶的容量
func RateLimitMiddleware(fillInterval time.Duration, cap int64) func(c *gin.Context) {
	bucket := ratelimit.NewBucket(fillInterval, cap)
	return func(c *gin.Context) {
		// 如果取不到令牌就中断本次请求返回 rate limit...
		/**
		bucket.TakeAvailable(count int64)
		内部调用tb.takeAvailable(tb.clock.Now(), count)
		当令牌桶中令牌数量为<=0时 返回0
		当请求的令牌=0时 返回0
		当请求的令牌数量 > 令牌桶中的令牌数量时 取出令牌桶中的令牌
		当请求的令牌数量 <= 令牌桶中的令牌数量时 令牌桶中的令牌数 - 请求的令牌数 取出请求的令牌
		*/
		//这个if表示 没取到令牌
		if bucket.TakeAvailable(1) < 1 {
			c.String(http.StatusOK, "rate limit...")
			c.Abort()
			return
		}
		//取到令牌 放行
		c.Next()
	}
}
